# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 11:09
# @Author  : Dell
# @File    : edit_question.py
# @Software: PyCharm
# @Desc    :问题改写：核心是# 多重查询：不同视角
# 流程：问题改写》并行搜索》知识整合》生成答案
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from operator import itemgetter

llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,
                 api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")

# --------------文档基础处理，索引和检索开始-----------------------
# 1、加载文本
loader = TextLoader(file_path="xiyouji.txt",encoding="utf-8")
blog_docs = loader.load()
# 2、分割文档
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=300,chunk_overlap=50)
splits = text_splitter.split_documents(blog_docs)
# 3、构建一个支持语义搜索的向量数据库（数据库在内存中）
embd=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
vectorstore = Chroma.from_documents(documents=splits,embedding=embd)
#4、检索
# 将一个向量数据库（vectorstore）转换为一个检索器（retriever），并设置检索参数：返回最相似的5个结果。
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
# --------------文档基础处理，索引和检索结束-----------------------

#---------------进行问题改写和多重检索开始-------------------------
template = """你是一名 AI 语言模型助手。你的任务是生成给定用户问题的五个不同版本，
以从向量数据库中检索相关文档。通过生成用户问题的多个视角，你
的目标是帮助用户克服基于距离的相似性搜索的一些限制。
提供这些以换行符分隔的备选问题。原始问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)
# 按提示词模版进行问题改写  ，生成5个问题
generate_queries = (
    prompt
    | llm
    | StrOutputParser()
    | (lambda x: x.split("\n"))
)
#多重检索答案的后期处理
def get_unique_union(documents:list[list]):
    """
    检索到的文档的唯一联合
    对嵌套的Document列表进行展平、去重，返回唯一的Document集合。

    Args:
        documents: 嵌套的Document列表（如 [[doc1, doc2], [doc3]]）

    Returns:
        list[Document]: 去重后的Document列表
    """
    # 展平嵌套列表
    flattened_docs = [doc for sublist in documents for doc in sublist]
    # 基于Document的page_content去重（假设page_content是唯一标识）
    unique_docs = []
    seen_contents = set()

    for doc in flattened_docs:
        if doc.page_content not in seen_contents:
            seen_contents.add(doc.page_content)
            unique_docs.append(doc)

    return unique_docs
# 检索
question = "根据这篇小说，你能介绍一下孙悟空的生活现状吗？"
"""核心关注的该链，问题改写》并行检索答案》答案去重后合并生成唯一结果"""
retrieval_chain = generate_queries | retriever.map() | get_unique_union
docs = retrieval_chain.invoke({"question":question})

#---------------进行问题改写和多重检索结束-------------------------
#---------------构建rag链进行问题回答-----------------------------
# 把所有都串起来，完整的多重问题改写RAG
template1 = """根据此上下文回答以下问题：

{context}

问题：{question}
"""

prompt1 = ChatPromptTemplate.from_template(template1)
# 构建最终rag链
final_rag_chain = (
    {"context": retrieval_chain,
     "question": itemgetter("question")}
    | prompt1
    | llm
    | StrOutputParser()
)

f_re=final_rag_chain.invoke({"question":question})
print(f"最后答复结果：{f_re}")

#---------------构建rag链进行问题回答结束-----------------------------